Skip to content

Conversation

@jessegrabowski
Copy link
Member

@jessegrabowski jessegrabowski commented Jan 2, 2026

Description

Split from #1806 . Implements gradients for join_dims and split_dims, now with 100% less reshape calls!

Also "improves" type hints for join_dims and split_dims.

Related Issue

  • Closes #
  • Related to #

Checklist

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

@jessegrabowski jessegrabowski added enhancement New feature or request gradients labels Jan 2, 2026
@jessegrabowski jessegrabowski marked this pull request as ready for review January 2, 2026 17:24
Improve type hints for `join_dims` and `split_dims`
@jessegrabowski jessegrabowski merged commit 0b439c0 into pymc-devs:main Jan 2, 2026
66 checks passed
@jessegrabowski jessegrabowski deleted the split-join-grad branch January 2, 2026 22:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request gradients

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants